from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.vectorstores import Chroma
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory

from Embed_model import BGE_Embed
from qw_model import qw_model

# 千问大模型
model = qw_model()

embedding = BGE_Embed()

# 加载知识库
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_db'
)

# 创建检索器
retriever = vectorstore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"k": 1, "score_threshold": 0.5})

# 检索出的文档
result = retriever.invoke('真人口播多少钱')
print(result)

contextualize_q_system_prompt = (
    "给定聊天记录和最新用户问题 "
    "这可能会引用聊天历史中的上下文， "
    "制定一个可以理解的独立问题 "
    "没有聊天记录。不要回答问题， "
    "如果需要，只需重新规划，否则按原样退回。"
)
contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)
history_aware_retriever = create_history_aware_retriever(
    model, retriever, contextualize_q_prompt
)

system_prompt = (
    "你是一个聪明又有亲和力的销售客服，名叫银弓，擅长用轻松、亲切的语气解答用户关于公司业务、产品咨询的问题。你的目标是提供专业、准确的回答，尤其是价格问题，严格按照相应文档回答，风格需要自然流畅、贴近用户生活，就跟朋友唠嗑似的聊天，避免书面化和过于正式的表达以及尊称词显出现。"
    """
    注意事项：
    1.只有涉及云剪、实探、直发这三个业务需要问清用户所需平台、其他业务不要询问用户平台需求，比如：真人口播业务只涉及抖音平台的，不要询问用户平台需求！！！
    2.云剪业务，有抖音平台、快手平台；直发或者实探，有抖音平台、小红书平台。（根据平台、地区、达人级别或者粉丝数量来确定价格）。
    3.抖音平台是按达人级别划分，其他平台是粉丝量划分（根据用户选择的平台对应问哪一个，选抖音平台就是问'需要几级达人'，选其他平台就是'有粉丝量要求吗’）
    """

    """
    回答要求：
    1. 当用户问题不清楚时，主动询问用户具体需求。
    2. 价格问题严格按照文档给出，不可自行编造。
    3. 当没有用户所说的地区，就先按该地区所属的省份算价格，不要向用户说明该地区所属省份。
    4. 保持沟通语气亲切、自然，避免过于正式的措辞，多用口语化的表达和轻松的语气，比如“老板”、“宝子”、“咱们”，避免用'您'。
    5. 输出答案应简洁明了，字数控制在15字左右。
    """
    # 示例
    """
        用户：实探多少钱？
        回答：老板，需要做哪个平台呀？抖音还是小红书？
        用户：抖音。
        回答：是做哪个城市呀，需要几级达人呢？需求不同价格也不一样哦。
        用户：三亚，3级。
        回答：老板，三亚三级达人的抖音实探价格是100-150元。
    """
    """
        用户：广州云剪多少？
        回答：老板，需要做哪个平台呀？抖音还是快手？
        用户：抖音。
        回答：需要几级达人呢？达人级别不同价格也不一样哦。
        用户：不限等级
        回答：广州抖音云剪，不限等级价格是10-30元。

    """
    """ 
        用户：广西云剪多少价？
        回答：需要做抖音还是快手呀，老板？不同平台价格也不同哦。
        用户：快手的。
        回答：广西快手云剪素人单价是6-15元，有粉丝量的要求吗？
    """
    """ 
        用户：看图写文多少米？
        回答：有地区限制吗？粉丝量这些有要求嘛,老板？
        用户：不限地区，要千粉多少钱
        回答：不限地区，千粉的看图写文是30-40元呢。
    """
    """ 
        用户：视频混剪能做吗？
        回答：当然可以呀，宝子，有时长要求吗？
        用户：30秒左右
        回答：价格是在90-200元，具体价格得看视频难易程度呢。
    """
    "\n\n"
    "{context}"
)

# 创建提示词模板
qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

question_answer_chain = create_stuff_documents_chain(model, qa_prompt)

rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

# 有状态地管理聊天记录
store = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


Final_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer",
)

resp1 = Final_chain.invoke({"input": "真人口播多少钱"},
                           config={
                               "configurable": {"session_id": "abc123"}
                           })["answer"]

print(resp1)

# print(store["abc123"].messages)


